package 堆;

import java.util.*;

/**
 * 373. 查找和最小的 K 对数字
 *
 * 给定两个以 非递减顺序排列 的整数数组 nums1 和 nums2 , 以及一个整数 k 。
 * 定义一对值 (u,v)，其中第一个元素来自 nums1，第二个元素来自 nums2 。
 * 请找到和最小的 k 个数对 (u1,v1),  (u2,v2)  ...  (uk,vk) 。
 *
 * 解决这个问题的一个高效方法是使用最小堆（优先队列）。具体步骤如下：
 *  1：初始化最小堆：首先，我们需要一个能够比较元素大小的最小堆。在Java中，可以通过PriorityQueue实现。
 *  2：初始化加入：将第一个数组的第一个元素与第二个数组的所有元素组成的对加入到堆中。这是因为对于每个后续步骤，我们总是基于当前已考虑过的最小对来生成新的候选对。
 *  3：迭代取出并更新：执行K次操作，每次从堆中取出和最小的一对，并根据这对中的元素在各自数组中的位置，尝试将下一个可能的对加入堆中。
 *  例如，如果从堆中取出的是(i, j)，则检查是否可以将(i+1, j)加入堆中（假设i+1在第一个数组的有效范围内）。
 *  4：结果收集：每次从堆中取出最小对时，将其加入结果列表中，直到获取了K个对或者堆为空为止。
 */
public class L_373 {

    /**
     * 优先队列（最小堆）：用于存储和排序可能的元素对，每次提取和最小的元素对。
     * 哈希集合：记录已访问的索引对，避免重复处理。
     * 初始索引对：从 (0, 0) 开始，将其加入堆中。
     * 循环提取最小对：每次从堆中取出和最小的索引对，将其结果添加到列表中，并扩展可能的下一个候选对（i+1, j）和 (i, j+1)，确保不会重复处理。
     * @param nums1
     * @param nums2
     * @param k
     * @return
     */
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        List<List<Integer>> result = new ArrayList<>();
        if (nums1 == null || nums1.length == 0 || nums2 == null || nums2.length == 0 || k <= 0) {
            return result;
        }

        int m = nums1.length;
        int n = nums2.length;

        // 使用优先队列（最小堆）
        PriorityQueue<int[]> minHeap = new PriorityQueue<>(
                (a, b) -> {
                    // 计算两个索引对对应的元素和的差值，作为比较的依据
                    long sumA = (long) nums1[a[0]] + nums2[a[1]];
                    long sumB = (long) nums1[b[0]] + nums2[b[1]];
                    return Long.compare(sumA, sumB);
                }
        );

        // 使用哈希集合来记录已经访问过的索引对，避免重复处理
        Set<Long> visited = new HashSet<>();
        int i = 0, j = 0;
        // 将初始索引对 (0, 0) 添加到堆中，并标记为已访问
        long key = ((long) i) << 32 | j; // 将i和j组合成一个64位整数，作为哈希集合的键
        minHeap.offer(new int[]{i, j});
        visited.add(key);

        // 循环，直到找到k对元素或者堆为空
        while ( k > 0 && !minHeap.isEmpty()){
            // 获取堆顶的索引对（即当前和最小的）
            int[] currentPair = minHeap.poll();
            i = currentPair[0];
            j = currentPair[1];

            // 将当前元素对添加到结果集中
            List<Integer> pair = new ArrayList<>();
            pair.add(nums1[i]);
            pair.add(nums2[j]);
            result.add(pair);

            // 搜索可能的下一个候选对（i+1 和 j）
            if (i + 1 < m){
                // 如果该索引对未被访问过，加入堆并标记为已访问
                key = ((long) (i + 1)) << 32 | j;
                if (!visited.contains(key)){
                    minHeap.offer(new int[]{i + 1, j});
                    visited.add(key);
                }
            }
            // 搜索可能的下一个候选对 （i 和 j+1）
            if (j + 1 < n){
                key = ((long) i) << 32 | (j + 1);
                if (!visited.contains(key)){
                    minHeap.offer(new int[]{i, j + 1});
                    visited.add(key);
                }
            }

            k--; // 添加一个元素后，减少剩余的k值
        }
        return result;
    }
}
